import sys
import unittest
import json
from unittest.mock import Mock
from evaluates.Metrics import MetricsComputer
from transformers import EvalPrediction

import os
# 添加项目根目录到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))


class TestMetricsComputer(unittest.TestCase):

    def setUp(self):
        # 创建一个模拟的tokenizer
        self.mock_tokenizer = Mock()
        self.mock_tokenizer.pad_token_id = 0
        self.mock_tokenizer.batch_decode = Mock()

        # 创建MetricsComputer实例
        self.metrics_computer = MetricsComputer(self.mock_tokenizer)

    def test_parse_entities_valid_json(self):
        # 测试解析有效的JSON实体
        text = 'some text ### Response: [{"entity": "entity1", "label": "label1"}, {"entity": "entity2", "label": "label2"}]'
        result = self.metrics_computer._parse_entities(text)
        expected = {("entity1", "label1"), ("entity2", "label2")}
        self.assertEqual(result, expected)

    def test_parse_entities_invalid_json(self):
        # 测试解析无效的JSON实体
        text = 'some text ### Response: invalid json'
        result = self.metrics_computer._parse_entities(text)
        self.assertEqual(result, set())

    def test_parse_entities_no_json(self):
        # 测试没有JSON的情况
        text = 'some text without JSON'
        result = self.metrics_computer._parse_entities(text)
        self.assertEqual(result, set())

    def test_parse_entities_missing_fields(self):
        # 测试实体缺少必要字段的情况
        text = '### Response: [{"entity": "entity1"}, {"label": "label2"}, {"entity": "entity3", "label": "label3"}]'
        result = self.metrics_computer._parse_entities(text)
        expected = {("entity3", "label3")}
        self.assertEqual(result, expected)

    def test_parse_entities_empty_array(self):
        # 测试空JSON数组的情况
        text = '### Response: []'
        result = self.metrics_computer._parse_entities(text)
        self.assertEqual(result, set())

    def test_parse_entities_malformed_json(self):
        # 测试格式错误的JSON
        text = '### Response: [{entity: "test", "label": "test"}]'  # 缺少引号
        result = self.metrics_computer._parse_entities(text)
        self.assertEqual(result, set())

    def test_parse_entities_empty_string(self):
        # 测试空字符串输入
        text = ''
        result = self.metrics_computer._parse_entities(text)
        self.assertEqual(result, set())

    def test_call_method_with_perfect_match(self):
        # 测试预测和标签完全匹配的情况
        # 模拟预测和标签
        preds = [[101, 202, 303]]
        labels = [[101, 202, 303]]

        # 设置tokenizer的batch_decode返回值
        self.mock_tokenizer.batch_decode.return_value = [
            '### Response: [{"entity": "entity1", "label": "label1"}]',
            '### Response: [{"entity": "entity1", "label": "label1"}]'
        ]

        eval_preds = EvalPrediction(predictions=preds, label_ids=labels)
        result = self.metrics_computer(eval_preds)

        # 当预测完全匹配时，precision, recall和f1都应该是1.0
        self.assertEqual(result["precision"], 1.0)
        self.assertEqual(result["recall"], 1.0)
        self.assertEqual(result["f1"], 1.0)

    def test_call_method_with_no_match(self):
        # 测试预测和标签完全不匹配的情况
        preds = [[101, 202, 303]]
        labels = [[101, 202, 303]]

        # 设置tokenizer的batch_decode返回值
        self.mock_tokenizer.batch_decode.return_value = [
            '### Response: [{"entity": "entity1", "label": "label1"}]',
            '### Response: [{"entity": "entity2", "label": "label2"}]'
        ]

        eval_preds = EvalPrediction(predictions=preds, label_ids=labels)
        result = self.metrics_computer(eval_preds)

        # 当预测完全不匹配时，precision, recall和f1都应该是0.0
        self.assertEqual(result["precision"], 0.0)
        self.assertEqual(result["recall"], 0.0)
        self.assertEqual(result["f1"], 0.0)

    def test_call_method_with_partial_match(self):
        # 测试预测和标签部分匹配的情况
        preds = [[101, 202, 303]]
        labels = [[101, 202, 303]]

        # 设置tokenizer的batch_decode返回值
        self.mock_tokenizer.batch_decode.return_value = [
            '### Response: [{"entity": "entity1", "label": "label1"}, {"entity": "entity2", "label": "label2"}]',
            '### Response: [{"entity": "entity1", "label": "label1"}, {"entity": "entity3", "label": "label3"}]'
        ]

        eval_preds = EvalPrediction(predictions=preds, label_ids=labels)
        result = self.metrics_computer(eval_preds)

        # 验证返回的是字典且包含所需键
        self.assertIn("precision", result)
        self.assertIn("recall", result)
        self.assertIn("f1", result)
        self.assertIsInstance(result["precision"], float)
        self.assertIsInstance(result["recall"], float)
        self.assertIsInstance(result["f1"], float)

    def test_call_method_empty_entities(self):
        # 测试没有实体的情况
        preds = [[101, 202, 303]]
        labels = [[101, 202, 303]]

        # 设置tokenizer的batch_decode返回值
        self.mock_tokenizer.batch_decode.return_value = [
            '### Response: []',
            '### Response: []'
        ]

        eval_preds = EvalPrediction(predictions=preds, label_ids=labels)
        result = self.metrics_computer(eval_preds)

        # 当没有实体时，所有指标都应该是0.0
        self.assertEqual(result["precision"], 0.0)
        self.assertEqual(result["recall"], 0.0)
        self.assertEqual(result["f1"], 0.0)

    def test_call_method_with_tuple_predictions(self):
        # 测试预测是元组的情况
        preds = ([[101, 202, 303]], None)  # 模拟元组形式的预测
        labels = [[101, 202, 303]]

        # 设置tokenizer的batch_decode返回值
        self.mock_tokenizer.batch_decode.return_value = [
            '### Response: [{"entity": "entity1", "label": "label1"}]',
            '### Response: [{"entity": "entity1", "label": "label1"}]'
        ]

        eval_preds = EvalPrediction(predictions=preds, label_ids=labels)
        result = self.metrics_computer(eval_preds)

        # 当预测完全匹配时，precision, recall和f1都应该是1.0
        self.assertEqual(result["precision"], 1.0)
        self.assertEqual(result["recall"], 1.0)
        self.assertEqual(result["f1"], 1.0)

    def test_call_method_with_labels_containing_neg100(self):
        # 测试标签包含-100的情况
        preds = [[101, 202, 303]]
        labels = [[101, -100, 303]]

        # 设置tokenizer的batch_decode返回值
        self.mock_tokenizer.batch_decode.return_value = [
            '### Response: [{"entity": "entity1", "label": "label1"}]',
            '### Response: [{"entity": "entity1", "label": "label1"}]'
        ]
        # 模拟将-100替换为pad_token_id的行为
        self.mock_tokenizer.pad_token_id = 0

        eval_preds = EvalPrediction(predictions=preds, label_ids=labels)
        # 验证不会抛出异常
        result = self.metrics_computer(eval_preds)
        self.assertIn("precision", result)
        self.assertIn("recall", result)
        self.assertIn("f1", result)

    def test_call_method_complex_scenario(self):
        # 测试更复杂的场景，多个实体和部分匹配
        preds = [[101, 202, 303, 404, 505]]
        labels = [[101, 202, 303, 404, 505]]

        # 设置tokenizer的batch_decode返回值，模拟多个响应
        self.mock_tokenizer.batch_decode.return_value = [
            '### Response: [{"entity": "entity1", "label": "label1"}, {"entity": "entity2", "label": "label2"}, {"entity": "entity3", "label": "label3"}]',
            '### Response: [{"entity": "entity1", "label": "label1"}, {"entity": "entity4", "label": "label4"}, {"entity": "entity3", "label": "label3"}]'
        ]

        eval_preds = EvalPrediction(predictions=preds, label_ids=labels)
        result = self.metrics_computer(eval_preds)

        # 验证返回的是字典且包含所需键
        self.assertIn("precision", result)
        self.assertIn("recall", result)
        self.assertIn("f1", result)
        self.assertIsInstance(result["precision"], float)
        self.assertIsInstance(result["recall"], float)
        self.assertIsInstance(result["f1"], float)


if __name__ == '__main__':
    unittest.main()
    #python -m unittest evaluates-test.py
